import numpy as np
import cv2
import os
import sys
import logging
sys.path.append(os.getcwd())
from api import Detector
from api import ImageOrientationCorrector
from api import ie
import time
labels = ['0','1','2','3','4','5','6','7','8','9']
colors = [[np.random.randint(0, 255) for _ in range(3)] for _ in range(len(labels))]
count = 0
def draw_dets(evaluator,inputs, dets_list, is_save=True, names=None):
    global count
    for img, dets, name in zip(inputs, dets_list, names):
        tl = round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
        tf = max(tl - 1, 1)  # font thickness
        if len(dets) > 0:
            if dets[0] is not None:
                for det in dets:
                    try:
                        cls, score, bbox = int(det[5]), det[4], det[:4]
                    except:
                        print()

                    c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))

                    patch = img[int(bbox[1]):int(bbox[3]),int(bbox[0]):int(bbox[2])]
                    probs = evaluator.run([patch])[0]
                    cls = np.argmax(probs)
                    label = labels[cls]
                    color = colors[cls]
                    t_size = cv2.getTextSize(label, 0, fontScale=tl / 4, thickness=tf)[0]

                    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)

                    c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
                    cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)
                    cv2.putText(img, label+' %.2f' % score, (c1[0], c1[1] - 2), 0, tl / 4, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
                    # print(name+' '+label)
            else:
                print(name+' does not have any target')
        if is_save:
            cv2.imwrite('../outputs/%s' %name, img)
        else:
            return img

def get_img(root):
    files = os.listdir(root)
    imgs = []
    file_names = []
    for file_name in files:
        img_path = os.path.join(root,file_name)
        img = cv2.imread(img_path)
        imgs.append(img)
        file_names.append(file_name)
    return imgs,file_names

def get_img(root):
    files = os.listdir(root)
    imgs = []
    file_names = []
    for file_name in files:
        img_path = os.path.join(root,file_name)
        imgs.append(img_path)
        file_names.append(file_name)
    return imgs,file_names


if __name__ == "__main__":  
    logging.basicConfig(level=logging.DEBUG)
    model_name = 'gridDet_0.1.0.onnx'
    det = Detector(model_path=os.path.join('../data',model_name),batch_size=25)
    evaluator = ImageOrientationCorrector(
        model_path='/home/lixuan/workspace/project/dectetion/shelf_grid_detection/data/model.onnx',
    )
    imgs,file_names = get_img('/home/lixuan/workspace/project/zhejiang/classify/samples')
    for index,img in enumerate(imgs):
        img = cv2.imread(img)
        outputs = det.detect([img])
        # print(outputs[0])
        draw_dets(evaluator,[img], outputs, is_save=True, names=[file_names[index]])